package com.kastrax.ai2db.connection.connector

import com.kastrax.ai2db.connection.model.*
import kotlinx.coroutines.Dispatchers
import kotlinx.coroutines.withContext
import org.springframework.jdbc.core.simple.JdbcClient
import org.springframework.jdbc.support.GeneratedKeyHolder
import org.springframework.jdbc.datasource.SingleConnectionDataSource
import org.springframework.stereotype.Component
import java.sql.Connection as JavaSqlConnection
import java.sql.Types
import java.sql.Timestamp
import java.sql.PreparedStatement
import java.time.Instant
import java.util.*
import javax.sql.DataSource

/**
 * MySQL implementation of the DatabaseConnector interface using JdbcClient
 */
@Component
class MySQLConnector : DatabaseConnector {
    
    /**
     * Connect to a MySQL database
     */
    override suspend fun connect(config: ConnectionConfig): Connection = withContext(Dispatchers.IO) {
        try {
            val dataSource = createDataSource(config)
            
            // Create JdbcClient from dataSource to test the connection
            val jdbcClient = JdbcClient.create(dataSource)
            
            // Validate connection
            jdbcClient.sql("SELECT 1").query(Int::class.java).single()
            
            return@withContext Connection(
                id = UUID.randomUUID().toString(),
                config = config,
                status = ConnectionStatus.CONNECTED,
                connectedAt = Instant.now(),
                dataSource = dataSource
            )
        } catch (e: Exception) {
            throw ConnectionException("Failed to connect to MySQL database: ${e.message}", e)
        }
    }
    
    /**
     * Disconnect from a MySQL database
     */
    override suspend fun disconnect(connection: Connection): Boolean = withContext(Dispatchers.IO) {
        try {
            (connection.dataSource as? SingleConnectionDataSource)?.destroy()
            return@withContext true
        } catch (e: Exception) {
            throw ConnectionException("Failed to disconnect from MySQL database: ${e.message}", e)
        }
    }
    
    /**
     * Test a MySQL database connection
     */
    override suspend fun testConnection(config: ConnectionConfig): ConnectionStatus = withContext(Dispatchers.IO) {
        try {
            val dataSource = createDataSource(config)
            val jdbcClient = JdbcClient.create(dataSource)
            
            // Test connection
            jdbcClient.sql("SELECT 1").query(Int::class.java).single()
            
            // Close connection
            (dataSource as? SingleConnectionDataSource)?.destroy()
            
            return@withContext ConnectionStatus.CONNECTED
        } catch (e: Exception) {
            return@withContext ConnectionStatus.FAILED
        }
    }
    
    /**
     * Get metadata about a MySQL database
     */
    override suspend fun getMetadata(connection: Connection): DatabaseMetadata = withContext(Dispatchers.IO) {
        try {
            val jdbcClient = JdbcClient.create(connection.dataSource!!)
            val tables = mutableListOf<TableMetadata>()
            val views = mutableListOf<TableMetadata>()
            val relationships = mutableListOf<Relationship>()

            // Get database metadata
            val metaData = connection.dataSource!!.connection.metaData
            val databaseProductName = metaData.databaseProductName
            val databaseProductVersion = metaData.databaseProductVersion
            val driverName = metaData.driverName
            val driverVersion = metaData.driverVersion

            // Get list of tables and views
            val tableRows = jdbcClient.sql(
                """
                SELECT TABLE_NAME, TABLE_TYPE, TABLE_COMMENT 
                FROM INFORMATION_SCHEMA.TABLES 
                WHERE TABLE_SCHEMA = :schema
                """
            )
            .param("schema", connection.config.database)
            .query { rs, _ ->
                mapOf(
                    "TABLE_NAME" to rs.getString("TABLE_NAME"),
                    "TABLE_TYPE" to rs.getString("TABLE_TYPE"),
                    "TABLE_COMMENT" to (rs.getString("TABLE_COMMENT") ?: "")
                )
            }
            .list()

            // Process each table
            for (tableRow in tableRows) {
                val tableName = tableRow["TABLE_NAME"] as String
                val tableType = tableRow["TABLE_TYPE"] as String
                val tableComment = tableRow["TABLE_COMMENT"] as String

                // Get columns for this table
                val columns = getTableColumns(jdbcClient, connection.config.database, tableName)

                // Get primary keys for this table
                val primaryKeys = getPrimaryKeys(jdbcClient, connection.config.database, tableName)

                // Get foreign keys for this table
                val foreignKeys = getForeignKeys(jdbcClient, connection.config.database, tableName)

                // Get indexes for this table
                val indexes = getIndexes(jdbcClient, connection.config.database, tableName)

                // Create table metadata
                val tableMetadata = TableMetadata(
                    name = tableName,
                    schema = connection.config.database,
                    type = tableType,
                    columns = columns,
                    primaryKeys = primaryKeys.map { it.columnNames }.flatten(),
                    foreignKeys = foreignKeys,
                    indexes = indexes,
                    comment = tableComment.ifEmpty { null }
                )

                // Add table or view
                if (tableType.contains("VIEW", ignoreCase = true)) {
                    views.add(tableMetadata)
                } else {
                    tables.add(tableMetadata)
                }

                // Identify relationships from foreign keys
                foreignKeys.forEach { foreignKey ->
                    val relationship = Relationship(
                        sourceTable = tableName,
                        sourceColumns = foreignKey.columnNames,
                        targetTable = foreignKey.referencedTableName,
                        targetColumns = foreignKey.referencedColumnNames,
                        type = RelationshipType.MANY_TO_ONE, // Default to MANY_TO_ONE
                        name = foreignKey.name,
                        foreignKeyName = foreignKey.name
                    )
                    relationships.add(relationship)
                }
            }

            // Analyze relationships to determine their types
            val analyzedRelationships = analyzeRelationships(relationships, tables)

            return@withContext DatabaseMetadata(
                databaseName = connection.config.database,
                schemaName = connection.config.database,
                tables = tables,
                views = views,
                databaseProductName = databaseProductName,
                databaseProductVersion = databaseProductVersion,
                driverName = driverName,
                driverVersion = driverVersion,
                relationships = analyzedRelationships,
                fetchedAt = Instant.now()
            )
        } catch (e: Exception) {
            throw QueryException("Failed to get database metadata: ${e.message}", e)
        }
    }

    /**
     * Get columns for a specific table
     */
    private fun getTableColumns(
        jdbcClient: JdbcClient,
        schema: String,
        tableName: String
    ): List<ColumnMetadata> {
        val columns = jdbcClient.sql(
            """
            SELECT 
                COLUMN_NAME, 
                DATA_TYPE, 
                COLUMN_TYPE,
                CHARACTER_MAXIMUM_LENGTH, 
                IS_NULLABLE, 
                COLUMN_KEY, 
                EXTRA, 
                COLUMN_DEFAULT, 
                COLUMN_COMMENT, 
                ORDINAL_POSITION 
            FROM INFORMATION_SCHEMA.COLUMNS 
            WHERE TABLE_SCHEMA = :schema AND TABLE_NAME = :tableName 
            ORDER BY ORDINAL_POSITION
            """
        )
        .param("schema", schema)
        .param("tableName", tableName)
        .query { rs, _ ->
            ColumnMetadata(
                name = rs.getString("COLUMN_NAME"),
                dataType = rs.getString("DATA_TYPE"),
                typeName = rs.getString("COLUMN_TYPE"),
                size = rs.getObject("CHARACTER_MAXIMUM_LENGTH")?.toString()?.toIntOrNull() ?: 0,
                nullable = rs.getString("IS_NULLABLE") == "YES",
                primaryKey = rs.getString("COLUMN_KEY") == "PRI",
                autoIncrement = rs.getString("EXTRA").contains("auto_increment", ignoreCase = true),
                defaultValue = rs.getString("COLUMN_DEFAULT"),
                comment = rs.getString("COLUMN_COMMENT").ifEmpty { null },
                ordinalPosition = rs.getInt("ORDINAL_POSITION")
            )
        }
        .list()

        return columns
    }

    /**
     * Get primary keys for a specific table
     */
    private fun getPrimaryKeys(
        jdbcClient: JdbcClient,
        schema: String,
        tableName: String
    ): List<Index> {
        val primaryKeys = jdbcClient.sql(
            """
            SELECT 
                k.CONSTRAINT_NAME, 
                k.COLUMN_NAME
            FROM INFORMATION_SCHEMA.TABLE_CONSTRAINTS t 
            JOIN INFORMATION_SCHEMA.KEY_COLUMN_USAGE k 
            ON t.CONSTRAINT_NAME = k.CONSTRAINT_NAME 
            WHERE 
                t.TABLE_SCHEMA = :schema 
                AND t.TABLE_NAME = :tableName 
                AND t.CONSTRAINT_TYPE = 'PRIMARY KEY'
            ORDER BY k.ORDINAL_POSITION
            """
        )
        .param("schema", schema)
        .param("tableName", tableName)
        .query { rs, _ ->
            mapOf(
                "CONSTRAINT_NAME" to rs.getString("CONSTRAINT_NAME"),
                "COLUMN_NAME" to rs.getString("COLUMN_NAME")
            )
        }
        .list()

        // Group by constraint name
        return primaryKeys
            .groupBy { it["CONSTRAINT_NAME"] as String }
            .map { (constraintName, columns) ->
                Index(
                    name = constraintName,
                    columnNames = columns.map { it["COLUMN_NAME"] as String },
                    unique = true,
                    type = "PRIMARY"
                )
            }
    }

    /**
     * Get foreign keys for a specific table
     */
    private fun getForeignKeys(
        jdbcClient: JdbcClient,
        schema: String,
        tableName: String
    ): List<ForeignKey> {
        val foreignKeys = jdbcClient.sql(
            """
            SELECT 
                k.CONSTRAINT_NAME, 
                k.COLUMN_NAME, 
                k.REFERENCED_TABLE_NAME, 
                k.REFERENCED_COLUMN_NAME,
                r.UPDATE_RULE,
                r.DELETE_RULE
            FROM INFORMATION_SCHEMA.KEY_COLUMN_USAGE k
            JOIN INFORMATION_SCHEMA.REFERENTIAL_CONSTRAINTS r
            ON k.CONSTRAINT_NAME = r.CONSTRAINT_NAME
            WHERE 
                k.TABLE_SCHEMA = :schema 
                AND k.TABLE_NAME = :tableName 
                AND k.REFERENCED_TABLE_NAME IS NOT NULL
            ORDER BY k.CONSTRAINT_NAME, k.ORDINAL_POSITION
            """
        )
        .param("schema", schema)
        .param("tableName", tableName)
        .query { rs, _ ->
            mapOf(
                "CONSTRAINT_NAME" to rs.getString("CONSTRAINT_NAME"),
                "COLUMN_NAME" to rs.getString("COLUMN_NAME"),
                "REFERENCED_TABLE_NAME" to rs.getString("REFERENCED_TABLE_NAME"),
                "REFERENCED_COLUMN_NAME" to rs.getString("REFERENCED_COLUMN_NAME"),
                "UPDATE_RULE" to rs.getString("UPDATE_RULE"),
                "DELETE_RULE" to rs.getString("DELETE_RULE")
            )
        }
        .list()

        // Group by constraint name
        return foreignKeys
            .groupBy { it["CONSTRAINT_NAME"] as String }
            .map { (constraintName, columns) ->
                ForeignKey(
                    name = constraintName,
                    columnNames = columns.map { it["COLUMN_NAME"] as String },
                    referencedTableName = columns.first()["REFERENCED_TABLE_NAME"] as String,
                    referencedColumnNames = columns.map { it["REFERENCED_COLUMN_NAME"] as String },
                    updateRule = columns.first()["UPDATE_RULE"] as String,
                    deleteRule = columns.first()["DELETE_RULE"] as String
                )
            }
    }

    /**
     * Get indexes for a specific table
     */
    private fun getIndexes(
        jdbcClient: JdbcClient,
        schema: String,
        tableName: String
    ): List<Index> {
        val indexes = jdbcClient.sql(
            """
            SELECT
                INDEX_NAME,
                COLUMN_NAME,
                NON_UNIQUE,
                INDEX_TYPE
            FROM INFORMATION_SCHEMA.STATISTICS
            WHERE
                TABLE_SCHEMA = :schema
                AND TABLE_NAME = :tableName
                AND INDEX_NAME != 'PRIMARY'
            ORDER BY INDEX_NAME, SEQ_IN_INDEX
            """
        )
        .param("schema", schema)
        .param("tableName", tableName)
        .query { rs, _ ->
            mapOf(
                "INDEX_NAME" to rs.getString("INDEX_NAME"),
                "COLUMN_NAME" to rs.getString("COLUMN_NAME"),
                "NON_UNIQUE" to rs.getBoolean("NON_UNIQUE"),
                "INDEX_TYPE" to rs.getString("INDEX_TYPE")
            )
        }
        .list()

        // Group by index name
        return indexes
            .groupBy { it["INDEX_NAME"] as String }
            .map { (indexName, columns) ->
                Index(
                    name = indexName,
                    columnNames = columns.map { it["COLUMN_NAME"] as String },
                    unique = !(columns.first()["NON_UNIQUE"] as Boolean),
                    type = columns.first()["INDEX_TYPE"] as String
                )
            }
    }

    /**
     * Analyze relationships to determine their types
     */
    private fun analyzeRelationships(
        relationships: List<Relationship>,
        tables: List<TableMetadata>
    ): List<Relationship> {
        // Create a map of tables by name for quick lookup
        val tableMap = tables.associateBy { it.name }

        // Create a map to count how many foreign keys point to each table
        val referenceCount = mutableMapOf<String, Int>()

        // Count references
        relationships.forEach { relationship ->
            referenceCount[relationship.targetTable] =
                (referenceCount[relationship.targetTable] ?: 0) + 1
        }

        // Analyze each relationship
        return relationships.map { relationship ->
            // Default type is MANY_TO_ONE
            var type = RelationshipType.MANY_TO_ONE

            // Get source and target tables
            val sourceTable = tableMap[relationship.sourceTable]
            val targetTable = tableMap[relationship.targetTable]

            if (sourceTable != null && targetTable != null) {
                // Check if source column is unique (primary key or unique index)
                val sourceColumnsAreUnique = sourceTable.columns
                    .filter { relationship.sourceColumns.contains(it.name) }
                    .all { it.primaryKey } ||
                    sourceTable.indexes.any { index ->
                        index.unique &&
                        relationship.sourceColumns.containsAll(index.columnNames) &&
                        index.columnNames.containsAll(relationship.sourceColumns)
                    }

                // Check if target column is unique
                val targetColumnsAreUnique = targetTable.columns
                    .filter { relationship.targetColumns.contains(it.name) }
                    .all { it.primaryKey } ||
                    targetTable.indexes.any { index ->
                        index.unique &&
                        relationship.targetColumns.containsAll(index.columnNames) &&
                        index.columnNames.containsAll(relationship.targetColumns)
                    }

                // Determine relationship type
                type = when {
                    sourceColumnsAreUnique && targetColumnsAreUnique -> RelationshipType.ONE_TO_ONE
                    sourceColumnsAreUnique && !targetColumnsAreUnique -> RelationshipType.ONE_TO_MANY
                    !sourceColumnsAreUnique && targetColumnsAreUnique -> RelationshipType.MANY_TO_ONE
                    else -> RelationshipType.MANY_TO_MANY
                }
            }

            // Return updated relationship
            relationship.copy(type = type)
        }
    }
    
    /**
     * Execute a query on a MySQL database
     */
    override suspend fun executeQuery(
        connection: Connection, 
        query: String,
        parameters: List<Any>,
        timeout: Long
    ): QueryResult = withContext(Dispatchers.IO) {
        val jdbcClient = JdbcClient.create(connection.dataSource!!)
        
        try {
            val startTime = System.currentTimeMillis()
            
            // Execute query with dynamic parameters
            val sqlOperation = jdbcClient.sql(query)
            
            // Add parameters if provided
            val parameterizedSql = if (parameters.isNotEmpty()) {
                // For positional parameters, add them in order
                parameters.foldIndexed(sqlOperation) { index, acc, param ->
                    acc.param(index + 1, param)
                }
            } else {
                sqlOperation
            }
            
            // Get column metadata
            val columns = mutableListOf<ColumnData>()
            // Store rows as Maps for flexibility
            val rows = mutableListOf<Map<String, Any?>>()

            // Execute query with automatic mapping to Map<String, Any?>
            parameterizedSql
                .query { rs, _ ->
                    // Extract column metadata if we haven't yet
                    if (columns.isEmpty()) {
                        val metaData = rs.metaData
                        val columnCount = metaData.columnCount

                        for (i in 1..columnCount) {
                            columns.add(
                                ColumnData(
                                    name = metaData.getColumnName(i),
                                    label = metaData.getColumnLabel(i),
                                    type = metaData.getColumnClassName(i),
                                    typeName = metaData.getColumnTypeName(i)
                                )
                            )
                        }
                    }

                    // Extract row data as Map
                    val row = mutableMapOf<String, Any?>()
                    for (column in columns) {
                        row[column.name] = rs.getObject(column.name)
                    }
                    rows.add(row)
                }
                .list()

            val executionTime = System.currentTimeMillis() - startTime
            
            return@withContext QueryResult(
                columns = columns,
                rows = rows,
                rowCount = rows.size,
                executionTime = executionTime,
                success = true
            )
        } catch (e: Exception) {
            // Return error result
            return@withContext QueryResult(
                error = e.message,
                success = false
            )
        }
    }
    
    /**
     * Execute an update operation on a MySQL database
     */
    override suspend fun executeUpdate(
        connection: Connection,
        query: String,
        parameters: List<Any>
    ): UpdateResult = withContext(Dispatchers.IO) {
        val jdbcClient = JdbcClient.create(connection.dataSource!!)
        
        try {
            val startTime = System.currentTimeMillis()
            val keyHolder = GeneratedKeyHolder()

            // Execute update and capture generated keys
            val updatedRows = connection.dataSource!!.connection.use { conn ->
                conn.prepareStatement(query, PreparedStatement.RETURN_GENERATED_KEYS).use { ps ->
                    // Set parameters
                    parameters.forEachIndexed { index, param ->
                        ps.setObject(index + 1, param)
                    }

                    // Execute update
                    ps.executeUpdate()

                    // Collect generated keys
                    val rs = ps.generatedKeys
                    val keys = mutableListOf<Any>()
                    while (rs.next()) {
                        keys.add(rs.getObject(1))
                    }

                    // Store keys in keyHolder
                    keys.forEach { keyHolder.keyList.add(mapOf("GENERATED_KEY" to it)) }

                    ps.updateCount
                }
            }

            val executionTime = System.currentTimeMillis() - startTime
            
            return@withContext UpdateResult(
                affectedRows = updatedRows,
                generatedKeys = keyHolder.keyList.map { it["GENERATED_KEY"] ?: it.values.first() },
                executionTime = executionTime,
                success = true
            )
        } catch (e: Exception) {
            return@withContext UpdateResult(
                error = e.message,
                success = false
            )
        }
    }
    
    /**
     * Begin a transaction
     */
    override suspend fun beginTransaction(connection: Connection): Transaction = withContext(Dispatchers.IO) {
        try {
            val jdbcConn = connection.dataSource!!.connection
            jdbcConn.autoCommit = false

            return@withContext Transaction(
                id = UUID.randomUUID().toString(),
                connection = connection,
                createdAt = Instant.now(),
                isActive = true
            )
        } catch (e: Exception) {
            throw TransactionException("Failed to begin transaction: ${e.message}", e)
        }
    }
    
    /**
     * Commit a transaction
     */
    override suspend fun commitTransaction(transaction: Transaction) = withContext(Dispatchers.IO) {
        try {
            val jdbcConn = transaction.connection.dataSource!!.connection
            jdbcConn.commit()
            jdbcConn.autoCommit = true
        } catch (e: Exception) {
            throw TransactionException("Failed to commit transaction: ${e.message}", e)
        }
    }
    
    /**
     * Rollback a transaction
     */
    override suspend fun rollbackTransaction(transaction: Transaction) = withContext(Dispatchers.IO) {
        try {
            val jdbcConn = transaction.connection.dataSource!!.connection
            jdbcConn.rollback()
            jdbcConn.autoCommit = true
        } catch (e: Exception) {
            throw TransactionException("Failed to rollback transaction: ${e.message}", e)
        }
    }
    
    /**
     * Create a DataSource for a MySQL connection
     */
    private fun createDataSource(config: ConnectionConfig): DataSource {
        val dataSource = SingleConnectionDataSource()
        dataSource.url = config.toJdbcUrl()
        dataSource.username = config.username
        dataSource.password = config.password
        dataSource.setSuppressClose(true)
        dataSource.setDriverClassName("com.mysql.cj.jdbc.Driver")
        return dataSource
    }
}

/**
 * Exception thrown when a connection operation fails
 */
class ConnectionException(message: String, cause: Throwable? = null) : RuntimeException(message, cause)

/**
 * Exception thrown when a query operation fails
 */
class QueryException(message: String, cause: Throwable? = null) : RuntimeException(message, cause)

/**
 * Exception thrown when a transaction operation fails
 */
class TransactionException(message: String, cause: Throwable? = null) : RuntimeException(message, cause)

